#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import re
import argparse
class kpoints():
    def __init__(self):
        self.nkpoints=100
        self.kpathlist=[]
        self.customK=0
        
    def readIBZKPT(self):
        ibzkpt=open('IBZKPT','r')
        kpoints=open('KPOINTS_0','w')
        count=0
        pline=''
        cline=''
        addKpoints=[]
        if self.kpathlist:
            print self.kpathlist
            NOcustomK=len(self.kpathlist)
            spanK=round(self.nkpoints/(NOcustomK-1))
            self.customK=spanK*(NOcustomK-1)+1
        for line in ibzkpt:
            if re.search('^\s+[0-9]+\s*$',line):
                originK=float(line.split()[0])
                if self.customK:
                    acturalK=originK+self.customK
                else:
                    spanK=round(self.nkpoints/(originK-1))
                    acturalK=spanK*(originK-1)+1+originK
                line='  '+str(int(acturalK))+'\n'

            if re.search('[0-9]\.[0-9]+\s+[0-9]\.[0-9]+\s+[0-9]\.[0-9]+\s+',line) and not self.customK:
                self.kpathlist.append(line)
                ###append kpath to
            kpoints.write(line)
                                                                                                                                                                                                                             
        for line in self.kpathlist:  
            if pline=='':
                pline=line
                addKpoints.append(str(float(pline.split()[0]))+'  '+str(float(pline.split()[1]))+'  '+str(float(pline.split()[2]))+'   0\n')
            else:
                cline=line
                clineX=float(cline.split()[0])
                clineY=float(cline.split()[1])
                clineZ=float(cline.split()[2])
                plineX=float(pline.split()[0])
                plineY=float(pline.split()[1])
                plineZ=float(pline.split()[2])
                for j in range(int(spanK)):
                    x=(clineX-plineX)/spanK*(j+1)+plineX
                    y=(clineY-plineY)/spanK*(j+1)+plineY
                    z=(clineZ-plineZ)/spanK*(j+1)+plineZ
                    addKpoints.append(str(x)+'  '+str(y)+'  '+str(z)+'   0\n')
                pline=cline
            
        for kp in addKpoints:
            kpoints.write(kp)
#            
parser=argparse.ArgumentParser(description='Generate the KPOINTS_0 file for band structure calculation in vasp. Need information from the IBZKPT file')
parser.add_argument("-n","--nkpoints", help="set the number of total K points, default:100",  type=int,default=100)
parser.add_argument("-p","--kpointPath", help="Set the custom k point path. format:\"0 0 0;0.5 0 0;0.5 0.5 0\" default path is read from IBZKPT file", type=str)  
#parser.add_argument("-k","--shrink",help="set the shrink(k points), e.g. '2,3,4,5,6 to benchmark different K points' Default: 2",  type=str,default='2') 
#parser.add_argument("-l","--levelshift",help="Set the level shift to help scf converge(1-10). Default: 0", type=str,default='0') 
#parser.add_argument("-m","--fmixing",help='Set the fmixing to help scf converge(1-50)%%. Default: 30', type=str,default='30')
#parser.add_argument("-d","--dimension",help='Set the dimension for periodicity (1,2,3). Default: 1', type=str,default='1')
#parser.add_argument("-s","--symmetry",help='Set the label for symmetry group. Default: 1', type=str,default='1')
#parser.add_argument("-o","--opt",help='Set the paramters for geometry optimization (ATOMONLY,FULLOPTG,CELLONLY,ITATOCEL,CVOLOPT). Default: FULLOPTG', nargs='?',type=str,default='False',const='FULLOPTG')
#parser.add_argument("-e","--external",help='Use external geometry(.gui) file as input', type=str)
#parser.add_argument("-g","--guess",help='Read initial guess from previous run', action='store_true')
#parser.add_argument("-D","--directory",help='Creat new directory with name filename-dirname',type=str)
#parser.add_argument("-v","--version",action='version', version='%(prog)s 2.0')
#parser.add_argument("structureFile",nargs='+')
args=parser.parse_args()            

kpoint=kpoints()

kpoint.nkpoints=args.nkpoints
if args.kpointPath:
    kpoint.kpathlist=args.kpointPath.split(';')
kpoint.readIBZKPT()
